#!/usr/bin/env python3
import uuid
import time
import threading
import queue
from concurrent import futures
import grpc

from cluster_fault_pb2 import Status, FaultMsgSignal, NodeFaultInfo, DeviceFaultInfo
from cluster_fault_pb2_grpc import FaultServicer, add_FaultServicer_to_server

from motor.utils.logger import get_logger

logger = get_logger(__name__)


class ClusterFaultServer(FaultServicer):
    def __init__(self):
        self.clients = {}  # job_id -> client_info
        self.fault_queue = queue.Queue()  # 故障消息队列
        self.running = True

    def Register(self, request, context):
        """注册客户端"""
        logger.info(f"Client registered: job_id={request.jobId}, role={request.role}")
        self.clients[request.jobId] = request
        return Status(code=0, info=f"Client {request.jobId}/{request.role} registered successfully")

    def SubscribeFaultMsgSignal(self, request, context):
        """订阅故障消息流"""
        logger.info(f"Client subscribed: job_id={request.jobId}, role={request.role}")

        try:
            while self.running and context.is_active():
                try:
                    # 阻塞等待故障消息，没有消息时会一直等待
                    fault_msg = self.fault_queue.get()
                    logger.info(f"Sending fault message: {fault_msg.uuid}")
                    yield fault_msg
                except Exception as e:
                    logger.error(f"Error getting fault message from queue: {e}")
                    break
        except Exception as e:
            logger.error(f"Error in SubscribeFaultMsgSignal: {e}")
        finally:
            logger.info(f"Client unsubscribed: job_id={request.jobId}")

    def add_fault_message(self, node_ip: str, fault_level: str, fault_codes: list = None,
                         device_type: str = "npu", device_id: str = "0"):
        """添加故障消息到队列"""
        if fault_codes is None:
            fault_codes = []

        # 创建设备故障信息
        device_fault = DeviceFaultInfo(
            deviceId=device_id,
            deviceType=device_type,
            faultCodes=fault_codes,
            faultLevel=fault_level,
            faultType=["hardware"],
            faultReason=["simulated fault"]
        )

        # 创建节点故障信息
        node_fault = NodeFaultInfo(
            nodeName=f"node-{node_ip.replace('.', '-')}",
            nodeIP=node_ip,
            nodeSN=f"SN-{node_ip}",
            faultLevel=fault_level,
            faultDevice=[device_fault]
        )

        # 创建故障消息信号
        fault_msg = FaultMsgSignal(
            uuid=str(uuid.uuid4()),
            jobId="mock-cluster",
            signalType="fault",
            nodeFaultInfo=[node_fault]
        )

        self.fault_queue.put(fault_msg)
        logger.info(f"Added fault message for node {node_ip} with level {fault_level}")


def create_server(host: str = "localhost", port: int = 5005):
    """创建gRPC服务器实例"""
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    fault_server = ClusterFaultServer()
    add_FaultServicer_to_server(fault_server, server)

    server_address = f"{host}:{port}"
    server.add_insecure_port(server_address)

    return server, fault_server, server_address


def run_server(host: str = "localhost", port: int = 5005):
    """运行gRPC服务器"""
    server, fault_server, server_address = create_server(host, port)
    server.start()

    logger.info(f"Cluster fault server started on {server_address}")

    try:
        server.wait_for_termination()
    except KeyboardInterrupt:
        logger.info("Server shutting down...")
        fault_server.running = False
        server.stop(0)

    return fault_server


if __name__ == "__main__":
    run_server()
